module parameter_setting

    implicit none

    ! model parameters

    ! parameters given by a matlab code ###############################################################3

    double precision :: parameters(13)
    double precision :: parameters_ss(16)
    double precision :: parameters_integer(8)
    double precision :: parameters_solution(2)
    double precision :: init_guess(6)
    integer :: init_type

    double precision :: beta
    double precision :: nu
    double precision :: alpha
    double precision :: delta
    double precision :: lambda
    double precision :: chi0
    double precision :: chi1
    double precision :: rho_a
    double precision :: sigma_a
    double precision :: nd

    double precision :: psi
    double precision :: n0
    double precision :: gamma
    double precision :: sigma_rk
    
    double precision :: c_ss,y_ss,i_ss,k_ss,h_ss,n_ss,l_ss,rb_ss,rkhats_ss,a_ss,xi_ss,pr_ss

    ! simulation
    double precision, allocatable :: av(:),kv(:),nv(:),rbv(:),xiv(:)
    double precision, allocatable :: cv(:),yv(:),iv(:),hv(:),lv(:),rhseev(:),pibv(:)
    double precision :: rkhat,rkhats,rkhat_next,rkhats_next
    double precision :: erk_next,pibv_next
    double precision :: xi_next,c_next,h_next,y_next,rhsee_next,rv_next,a_next,rb_next
    double precision :: z_temp
    double precision :: z_temp_log, cdf_log
    
    double precision, allocatable :: x0(:,:),x(:),x0_next(:,:),x_next(:),x0_reg(:,:),x_reg(:)
    double precision, allocatable :: rhsee_each(:),rhsee_derived(:)
    double precision :: dummy1,dummy2,rhsee_temp,cdf,pdf

    ! polynomials
    integer :: nstv,degp,npol ! number of state variables, degree of polynomial, size of polynomial matrix
    double precision, allocatable :: expmx(:,:),expmx_column(:),state(:)
    
    double precision, allocatable :: eta1_rhsee(:),eta1_rhsee_new(:,:),eta1_rhsee_old(:,:)
    
    integer :: con_max(1)
    double precision :: con_max_new(1,1), con_max_old(1,1)
    
    ! trapezoid integration
    integer :: n_trapz
    double precision, allocatable :: z_trapz(:)

    ! policy function iteration and convergence
    integer :: npf
    double precision :: tol, adj
    double precision, allocatable :: con_lev(:), rsq(:), shocks(:)
    integer :: t_max, iter_max, t_drop
    integer :: t, i, j, f, g, z, count

    ! values for non-linear solver
    integer :: nnls
    double precision, parameter :: tol_fun = 1d-8   
    double precision :: guess(1),fval(1)
!    double precision, allocatable :: guess(:),fval(:)
    integer :: info
    integer :: nonsol

    ! regression to update polynomial coefficients
    integer :: num_norun
    double precision, allocatable :: x_norun(:,:),xx_norun(:,:),xx_inv_norun(:,:)
    double precision, allocatable :: y_rhsee_norun(:,:),xy_rhsee_norun(:,:)
    double precision, allocatable :: y1_temp_norun(:,:),x_temp_norun(:,:)

    ! to import data
    character(len=20) :: filename               

end module

! main program ##########################################################################################   

program baseline_model

    ! modules
    use parameter_setting
    use my_toolbox

    implicit none
    
    ! import parameters from text files #################################################################

    open (10, file='from_matlab_to_fortran/parameters.txt', status='old', action='read')
       read(10,*) parameters
    close(10)
    beta     = parameters(4)
    nu       = parameters(5)
    alpha    = parameters(6)
    delta    = parameters(7)
    lambda   = parameters(8)
    chi0     = parameters(9)
    chi1     = parameters(10)
    rho_a    = parameters(11)
    sigma_a  = parameters(12)
    nd       = parameters(13)

    open (10, file='from_matlab_to_fortran/ss_values.txt', status='old', action='read')
       read(10,*) parameters_ss
    close(10)
    c_ss      = parameters_ss(1)
    y_ss      = parameters_ss(2)
    i_ss      = parameters_ss(3)
    k_ss      = parameters_ss(4)
    h_ss      = parameters_ss(5)
    n_ss      = parameters_ss(6)
    l_ss      = parameters_ss(7)
    rb_ss     = parameters_ss(8)
    rkhats_ss = parameters_ss(9)
    a_ss      = parameters_ss(10)
    xi_ss     = parameters_ss(11)
    pr_ss     = parameters_ss(12)
    psi       = parameters_ss(13)
    n0        = parameters_ss(14)
    gamma     = parameters_ss(15)
    sigma_rk  = parameters_ss(16)

    open (11, file='from_matlab_to_fortran/parameters_integer.txt', status='old', action='read')
        read(11,*) parameters_integer
    close(11)
    npf      = nint(parameters_integer(1)) ! number of policy functions to be approximated
    nstv     = nint(parameters_integer(2)) ! number of state variables
    degp     = nint(parameters_integer(3)) ! degree of polynomial
    npol     = nint(parameters_integer(4)) ! size of polynomial matrix
    iter_max = nint(parameters_integer(5))
    n_trapz  = nint(parameters_integer(6)) ! number of grids for integration approximation
    t_max    = nint(parameters_integer(7))
    t_drop   = nint(parameters_integer(8))

    open (11, file='from_matlab_to_fortran/parameters_solution.txt', status='old', action='read')
        read(11,*) parameters_solution
    close(11)
    tol = parameters_solution(1)
    adj = parameters_solution(2)

    allocate(expmx(nstv,npol))
    allocate(expmx_column(npol*nstv))
    open (11, file='from_matlab_to_fortran/polynomial_matrix.txt', status='old', action='read')
        read(11,*) expmx_column
    close(11)
    f=0
    do i=1,nstv
        do j=1,npol
            f=f+1
            expmx(i,j) = expmx_column(f)
        end do
    end do

    allocate(z_trapz(n_trapz))
    open (11, file='from_matlab_to_fortran/z_trapz.txt', status='old', action='read')
        read(11,*) z_trapz
    close(11)

    allocate(shocks(t_max+1))
    open (11, file='from_matlab_to_fortran/shocks.txt', status='old', action='read')
        read(11,*) shocks
    close(11)

    allocate(eta1_rhsee(npol),eta1_rhsee_new(npol,1),eta1_rhsee_old(npol,1))

    open (11, file='from_matlab_to_fortran/eta1_rhsee.txt', status='old', action='read')
        read(11,*) eta1_rhsee
    close(11)

    allocate(av(t_max+1),nv(t_max+1),kv(t_max+1),rbv(t_max+1),xiv(t_max+1))
    allocate(cv(t_max+1),yv(t_max+1),iv(t_max+1),lv(t_max+1),hv(t_max+1),rhseev(t_max+1),pibv(t_max+1))

    allocate(x_norun(t_max,npol),xx_inv_norun(npol,npol))    
    allocate(y_rhsee_norun(t_max,1),xy_rhsee_norun(npol,1))

    allocate(y1_temp_norun(t_max,1),x_temp_norun(t_max,npol))

    allocate(state(nstv))
    allocate(x0(nstv,npol),x(npol),x0_next(nstv,npol),x_next(npol),x0_reg(nstv,npol),x_reg(npol))
    
    allocate(rhsee_each(n_trapz))
    allocate(rhsee_derived(t_max))

    ! iterate until policy function converges ########################################################

    count = 1
!    do count = 1, iter_max

        ! set initial states
        av(1)  = a_ss
        kv(1)  = k_ss

        ! simulate for t_max periods
        do t = 2, t_max

        ! Variables determined independent of bank runs
        av(t)  = rho_a*av(t-1) + sigma_a*shocks(t-1)        
        hv(t)  = ((1d0-alpha)/psi*exp(av(t))*kv(t-1)**(alpha))**(1d0/(alpha+1d0/nu))
        yv(t)  = exp(av(t)) * kv(t-1)**(alpha)*hv(t)**(1-alpha)
        
        
        state(1) = log(kv(t-1))
        state(2) = av(t)
        do i=1,nstv
            do j=1,npol
                x0(i,j) = state(i)**expmx(i,j)
            end do
        end do
        x  = x0(1,:)*x0(2,:) ! *x0(3,:)*x0(4,:)
        
        rhseev(t) = exp( sum(x*eta1_rhsee) )
        cv(t) = rhseev(t)**(-1d0) + psi*hv(t)**(1d0+1d0/nu)/(1d0+1d0/nu)
        
        iv(t)  = yv(t) - cv(t)
        kv(t)  = (1d0-delta)*kv(t-1) + iv(t)

        rbv(t) = 1d0 - delta + alpha*yv(t)/kv(t-1)
                
        ! Numerical integration #########################################################################
        ! Compute right-hand side of Euler equation using numerical integration
        do z=1,n_trapz
        
            ! For each loop, take the same steps as above to compute endogenous variables.
            ! Each loop corresponds to a different a shock.
            
            a_next  = rho_a*av(t) + sigma_a*z_trapz(z)
            h_next  = ((1d0-alpha)/psi*exp(a_next)*kv(t)**(alpha))**(1d0/(alpha+1d0/nu))
            y_next  = exp(a_next) * kv(t)**(alpha)*h_next**(1-alpha)
                    
            state(1) = log(kv(t))
            state(4) = a_next
            do i=1,nstv
                do j=1,npol
                    x0_next(i,j) = state(i)**expmx(i,j)
                end do
            end do
            x_next  = x0_next(1,:)*x0_next(2,:) ! *x0_next(3,:)*x0_next(4,:)

            rhsee_next = exp( sum(x_next*eta1_rhsee) )
            c_next     = rhsee_next**(-1d0) + psi*h_next**(1d0+1d0/nu)/(1d0+1d0/nu)
            rb_next    = 1d0 - delta + alpha*y_next/kv(t)
            rhsee_temp = beta / ( c_next - psi*h_next**(1d0+1d0/nu)/(1d0+1d0/nu) ) * rb_next

            ! multiply density, and cut off the both tails outside of z_trapz
            call normp( z_trapz(z), dummy1, dummy2, pdf )    ! obtain pdf at z_trapz(z)
            call normp( z_trapz(1), cdf, dummy1, dummy2 )    ! obtain cdf of the lower tail
            rhsee_each(z) = rhsee_temp * pdf / ( 1d0 - 2d0*cdf )

        end do
        
        ! compute integeral
        call trapz_integ(z_trapz,rhsee_each,rhsee_derived(t)) ! rhsee_derived(2) corresponds to E2(uc(3))
        
        ! End of integration ############################################################################
        
!        write(*,*) t,'kv=',kv(t),'nv=',nv(t),'Rb=',rbv(t),'xis=',xisv(t)
        
        end do
        
        ! End of simulation #############################################################################
        
    ! Create variables in OLS    
    ! Be careful about timing of variables.
    ! Regress Et(uc(t+1)) on states k(t-1),n(t-1),...
    
    ! log of expected RHSEE starts from E1(uc(2)), which is not used.
    
    num_norun=0
    do t = t_drop+2, t_max

        state(1) = log(kv(t-1))
        state(2) = av(t)
        do i=1,nstv
            do j=1,npol
                x0_reg(i,j) = state(i)**expmx(i,j)
            end do
        end do
        x_reg  = x0_reg(1,:)*x0_reg(2,:) ! *x0_reg(3,:)*x0_reg(4,:)
        num_norun = num_norun + 1
        y1_temp_norun(num_norun,1) = log(rhsee_derived(t))
        x_temp_norun(num_norun,:)  = x_reg

    end do
        
    ! resize regression matrices
    call resize_2_real(y_rhsee_norun, num_norun, 1)    
    call resize_2_real(x_norun, num_norun, npol)
    
    y_rhsee_norun(:,1) = y1_temp_norun(1:num_norun,1)
    x_norun(:,:)       = x_temp_norun (1:num_norun,:)

    ! export derived rhsee,rbar,x

    open (50, file='from_fortran_to_matlab/rhsee_norun.txt', status='replace', action='write')
        do j=1,num_norun
            write(50,*) y_rhsee_norun(j,1)
        end do
    close(50)

    open (50, file='from_fortran_to_matlab/x_norun.txt', status='replace', action='write')
    do j=1,npol
        do i=1,num_norun
            write(50,*) x_norun(i,j)
        end do
    end do
    close(50)

    open (50, file='from_fortran_to_matlab/temp_eta1_rhsee.txt', status='replace', action='write')
    do i=1,npol
        write(50,*) eta1_rhsee(i)
    end do
    close(50)

    ! End of one simulation

end program
